import aiohttp
import json
import requests

from config.prompts_version3 import prompt_project_describe_version3
from network.beans import ResponseResult
from shared_variables import CONFIG

"""
    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": "将「太阳」翻译成意大利语"
            },
            {
                "role": "assistant",
                "content": "Il Sole"
            }
        ],
        "temperature": 0.8,
        "penalty_score": 2,
        "system": "你是一个大师"
    })

"""


class BaiduConnector:
    def __init__(self):
        self.API_KEY = CONFIG['api']['api_key']
        self.SECRET_KEY = CONFIG['api']['secret_key']

    async def send(self, message, temperature, penalty_score, system):
        url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + self.get_access_token()
        payload = json.dumps({
            "messages": message,
            "temperature": temperature,
            "penalty_score": penalty_score,
            "system": system
        })
        headers = {
            'Content-Type': 'application/json'
        }

        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(url, headers=headers, data=payload) as response:
                    response.raise_for_status()  # 检查请求是否成功，如果不成功会抛出异常
                    return ResponseResult(success=True, data=await response.text(), error_message=None)
        except aiohttp.ClientError as e:
            error_message = str(e)
            return ResponseResult(success=False, data=None, error_message=error_message)

    def send_sync(self, message, temperature, penalty_score, system):
        url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + self.get_access_token()
        payload = json.dumps({
            "messages": message,
            "temperature": temperature,
            "penalty_score": penalty_score,
            "system": system
        })
        headers = {
            'Content-Type': 'application/json'
        }

        try:
            response = requests.post(url, headers=headers, data=payload)
            response.raise_for_status()  # 检查请求是否成功，如果不成功会抛出异常
            return ResponseResult(success=True, data=response.text, error_message=None)
        except requests.exceptions.RequestException as e:
            error_message = str(e)
            return ResponseResult(success=False, data=None, error_message=error_message)

    def send_knowledge(self, message, input_variables={}):
        url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/plugin/rtnuih7emjryq3uf/?access_token=" + self.get_access_token()
        if input_variables:
            payload = json.dumps({
                "query": message,
                "plugins": ["uuid-zhishiku"],
                "input_variables": input_variables,
                "llm": {"temperature": 0.70, "penalty_score": 1}
            })
        else:
            payload = json.dumps({
                "query": message,
                "plugins": ["uuid-zhishiku"],
                "llm": {"temperature": 0.70, "penalty_score": 1}
            })
        headers = {
            'Content-Type': 'application/json'
        }

        try:
            response = requests.post(url, headers=headers, data=payload)
            response.raise_for_status()  # 检查请求是否成功，如果不成功会抛出异常
            return ResponseResult(success=True, data=response.text, error_message=None)
        except requests.exceptions.RequestException as e:
            error_message = str(e)
            return ResponseResult(success=False, data=None, error_message=error_message)

    def art(self, content):
        url = "https://aip.baidubce.com/rpc/2.0/ernievilg/v1/txt2img?access_token=" + self.get_access_token()
        payload = json.dumps({
            "text": content,
            "resolution": "1024*1536",
            "style": "二次元",
            "num": 1
        })
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        }

        try:
            response = requests.post(url, headers=headers, data=payload)
            response.raise_for_status()  # 检查请求是否成功，如果不成功会抛出异常
            return ResponseResult(success=True, data=response.text, error_message=None)
        except requests.exceptions.RequestException as e:
            error_message = str(e)
            return ResponseResult(success=False, data=None, error_message=error_message)

    def send_longer(self, message, temperature, penalty_score, system):
        url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie_bot_8k?access_token=" + self.get_access_token()
        payload = json.dumps({
            "messages": message,
            "temperature": temperature,
            "penalty_score": penalty_score,
            "system": system
        })
        headers = {
            'Content-Type': 'application/json'
        }

        try:
            response = requests.post(url, headers=headers, data=payload)
            response.raise_for_status()  # 检查请求是否成功，如果不成功会抛出异常
            return ResponseResult(success=True, data=response.text, error_message=None)
        except requests.exceptions.RequestException as e:
            error_message = str(e)
            return ResponseResult(success=False, data=None, error_message=error_message)

    def send_stream(self, message, temperature, penalty_score, system):
        url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + self.get_access_token()
        payload = json.dumps({
            "messages": message,
            "temperature": temperature,
            "penalty_score": penalty_score,
            "system": system,
            "stream": True
        })
        headers = {
            'Content-Type': 'application/json'
        }

        try:
            response = requests.post(url, headers=headers, data=payload, stream=True)  # 使用stream=True启用流式传输
            response.raise_for_status()

            # 通过迭代器逐行读取响应内容
            for line in response.iter_lines(decode_unicode=True):
                if line:
                    yield ResponseResult(success=True, data=line, error_message=None)

        except requests.exceptions.RequestException as e:
            error_message = str(e)
            yield ResponseResult(success=False, data=None, error_message=error_message)

    def get_access_token(self):
        """
        使用 AK，SK 生成鉴权签名（Access Token）
        :return: access_token，或是None(如果错误)
        """
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {"grant_type": "client_credentials", "client_id": self.API_KEY, "client_secret": self.SECRET_KEY}
        return str(requests.post(url, params=params).json().get("access_token"))


def test_zhishiku():
    client = BaiduConnector()

    temp_dict = {'boss': "曹操", 'me': "杨幂", 'project': "建造太空站", 'task': "端好茶"}

    response = client.send_knowledge("我需要做什么呢？", input_variables=temp_dict)
    print(response)
    message = prompt_project_describe_version3.format(temp_dict['boss'], temp_dict['me'], temp_dict['project'], temp_dict['me'], temp_dict['task'])
    print(message)
    response = client.send_knowledge(message, input_variables=temp_dict)
    print(response)

    # 这个回复可真是随机哦
    message = prompt_project_describe_version3.format(temp_dict['boss'], temp_dict['me'], temp_dict['project'], temp_dict['me'], temp_dict['task'])
    response = client.send_knowledge(message, input_variables={})
    print(response)

if __name__ == '__main__':

    client = BaiduConnector()
    response = client.art("生成一张照片，里面的内容是赤壁之战的前夕，杨幂正在求东风,可爱风格")
    print(response)